import torch.nn as nn

import torch.nn as nn


class Adapter(nn.Module):
    """
    Adapter 模块，用于通过降维和升维操作在神经网络中进行适配。实现了一个类似于残差连接（skip connection）的结构，
    使得网络在学习时更容易保持信息流通。
    """

    def __init__(self, dim, reduction_factor=16):
        """
        初始化 Adapter 模块。

        参数:
            dim (int): 输入张量的特征维度。
            reduction_factor (int): 降维因子，默认为16。用于缩小输入特征维度的比例。
        """
        super(Adapter, self).__init__()

        # 降维操作：将输入特征维度降到原来的 1/reduction_factor
        self.down_proj = nn.Linear(dim, dim // reduction_factor)
        # ReLU 激活函数，用于非线性变换
        self.relu = nn.ReLU()
        # 升维操作：将特征维度从降维后的维度恢复到原始维度
        self.up_proj = nn.Linear(dim // reduction_factor, dim)

    def forward(self, x):
        """
        前向传播函数。执行降维、ReLU 激活和升维操作，并返回带有残差连接的输出。

        参数:
            x (Tensor): 输入张量，形状为 (batch_size, dim)。

        返回:
            Tensor: 输出张量，形状与输入张量相同，经过适配模块处理后的结果。
        """
        residual = x  # 保存输入张量作为残差

        # 降维操作
        x = self.down_proj(x)
        # 应用 ReLU 激活函数
        x = self.relu(x)
        # 升维操作
        x = self.up_proj(x)

        # 返回输出张量，并添加残差连接
        return x + residual  # Skip connection（残差连接）

